import json
import re
from typing import List, Dict

# ====================== 配置参数 ======================
INPUT_JSON_PATH = "train_cf/data/大西瓜籽妈妈.deepseek.chat.cf.json"  # 替换为您的JSON文件路径
OUTPUT_JSONL_PATH = "train_cf/src/20250625/family_sft_data.jsonl"  # 输出的JSON Lines文件路径
INSTRUCTION_TEMPLATES = {
    "家庭成长背景类": "请总结于小戈的家庭成长背景（如出身、父母职业、成长环境）。",
    "家庭成员互动类": "请分享于小戈与家人的一个具体生活片段（如童年回忆、日常互动）。",
    "家庭对其职业的影响类": "家庭经历如何影响了于小戈的职业选择或价值观？",
    "家庭特殊事件类": "于小戈的家庭是否经历过对她影响深远的特殊事件？",
    "家庭与公众形象的关联类": "于小戈在公开场合如何提及家庭？反映了她怎样的家庭观念？"
}

# ====================== 核心函数 ======================
def extract_family_info(article: str) -> Dict:
    """从文章中提取家庭经历关键信息（基于规则匹配，可根据实际文章结构调整）"""
    info = {}
    
    # 家庭成长背景（匹配出身、父母职业、成长环境关键词）
    background_pattern = r"(出生于|成长于|家庭出身).*?(普通|教师|纺织工|商人等).*?(环境|氛围|教育)"
    background_match = re.search(background_pattern, article, re.DOTALL)
    if background_match:
        info["家庭成长背景"] = background_match.group(0)
    
    # 家庭成员互动（匹配具体生活场景关键词）
    interaction_pattern = r"(童年时|小时候|周末|假期).*?(母亲|父亲|阿崽｜籽籽｜前夫｜老公).*?(带我|教我|一起)(逛|学|做)(裁缝铺|旧书店|商场|厨房等)"
    interaction_match = re.search(interaction_pattern, article, re.DOTALL)
    if interaction_match:
        info["家庭成员互动"] = interaction_match.group(0)
    
    # 家庭对职业的影响（匹配职业相关影响关键词）
    career_impact_pattern = r"(影响|促使|坚定了|启蒙了).*?(时尚|服装|设计|品牌)行业的兴趣|选择"
    career_impact_match = re.search(career_impact_pattern, article, re.DOTALL)
    if career_impact_match:
        info["家庭对职业的影响"] = career_impact_match.group(0)
    
    # 家庭特殊事件（匹配搬家、重要活动等关键词）
    special_event_pattern = r"(搬家|迁移|随父母|重要|影响深远).*?(上海|北京|大城市|新环境)"
    special_event_match = re.search(special_event_pattern, article, re.DOTALL)
    if special_event_match:
        info["家庭特殊事件"] = special_event_match.group(0)
    
    return info

def generate_sft_samples(article: str) -> List[Dict]:
    """根据文章内容生成SFT指令-输出样本"""
    samples = []
    family_info = extract_family_info(article)
    
    # 为每类指令生成样本（若信息存在）
    for instr_type, instr_text in INSTRUCTION_TEMPLATES.items():
        if instr_type in family_info:
            sample = {
                "instruction": instr_text,
                "input": "无",  # 若需输入上下文可改为文章片段，此处简化为"无"
                "output": family_info[instr_type]
            }
            samples.append(sample)
        else:
            # 若信息不存在，生成提示信息（可选）
            sample = {
                "instruction": instr_text,
                "input": "无",
                "output": "文章中未明确提及与此相关的家庭经历。"
            }
            samples.append(sample)
    
    return samples

# ====================== 主流程 ======================
if __name__ == "__main__":
    # 读取原始JSON数据
    with open(INPUT_JSON_PATH, "r", encoding="utf-8") as f:
        raw_data = json.load(f)
    
    # 提取所有"output"字段的文章内容
    articles = [item.get("output", "") for item in raw_data if "output" in item]
    
    # 生成SFT数据集
    sft_samples = []
    for idx, article in enumerate(articles):
        if article.strip():  # 跳过空内容
            samples = generate_sft_samples(article)
            sft_samples.extend(samples)
            print(f"处理第 {idx+1} 篇文章，生成 {len(samples)} 条样本")
    
    # 保存为JSON Lines格式
    with open(OUTPUT_JSONL_PATH, "w", encoding="utf-8") as f:
        for sample in sft_samples:
            json.dump(sample, f, ensure_ascii=False, indent=2)
            f.write("")
    
    print(f"数据集生成完成！共生成 {len(sft_samples)} 条样本，保存至 {OUTPUT_JSONL_PATH}")